{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using PyMilvus's Model To Generate Text Embeddings\n",
    "Quickly enhance your search capabilities with text embeddings using PyMilvus(higher than 2.4.0). This guide shows how to utilize PyMilvus models for extracting rich text embeddings, setting the foundation for powerful search functionalities.  \n",
    "In this doc, we will go through **dense embedding** models, **sparse embedding** models, and **hybrid** models to show how to use them in action.  \n",
    "Let's begin by installing the dependencies(you could also use `virtualenv` to create a new enviroment):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install pymilvus[model]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For most use cases, you can generate embeddings for storage or retrieval by simply using `ef(texts)`. However, when you need different processing for queries and documents, you could use two specific functions. Documents are processed with `encode_documents` to generate their embeddings, which are then stored in the vector database. For retrieval, the query is processed using `encode_queries` to create its embedding, which is then used to search the database."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Dense embedding** is a technique used in natural language processing to represent words or phrases as continuous, dense vectors in a high-dimensional space, capturing semantic relationships."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## OpenAI Embedding Function\n",
    "OpenAI offers dense [embedding services](https://platform.openai.com/docs/guides/embeddings), but to access them, users must sign up and obtain an API key. With the API key properly set in your environment variables, you can start using tools like OpenAIEmbeddingFunction to generate dense embeddings.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dense dim: 512 (512,)\n",
      "dense dim: 512 (512,)\n"
     ]
    }
   ],
   "source": [
    "docs = [\n",
    "    \"Artificial intelligence was founded as an academic discipline in 1956.\",\n",
    "    \"Alan Turing was the first person to conduct substantial research in AI.\",\n",
    "    \"Born in Maida Vale, London, Turing was raised in southern England.\",\n",
    "]\n",
    "\n",
    "from pymilvus import model\n",
    "\n",
    "# initialize using 'text-embedding-3-large'\n",
    "openai_ef = model.dense.OpenAIEmbeddingFunction(\n",
    "    model_name=\"text-embedding-3-large\", # Specify the model name\n",
    "    dimensions=512 # Set the embedding dimensionality according to MRL feature.\n",
    ")\n",
    "\n",
    "# get the embeddings in general way.\n",
    "queries = docs\n",
    "queries_embeddings = openai_ef(queries)\n",
    "docs_embeddings = openai_ef(docs)\n",
    "\n",
    "# get the embeddings in specified way.\n",
    "queries_embeddings = openai_ef.encode_queries(queries)\n",
    "docs_embeddings = openai_ef.encode_documents(docs)\n",
    "\n",
    "# now we can check the dimension of embedding from results and the embedding function.\n",
    "print(\"dense dim:\", openai_ef.dim, queries_embeddings[0].shape)\n",
    "print(\"dense dim:\", openai_ef.dim, docs_embeddings[0].shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "when using OpenAIEmbeddingFunction, `encoding_queries` and `encoding_documents` are exactly same procedure, one can use `openai_ef(texts)` instead.  \n",
    "- `openai_ef(texts)`: same with other two functions.\n",
    "- `openai_ef.encode_queries(queries)`: same with other two functions.\n",
    "- `openai_ef.encode_documents(documents)`: same with other two functions.\n",
    "\n",
    "\n",
    "Additionally, you may initialize the OpenAIEmbeddingFunction by directly providing OpenAI's official parameters, such as api_key and base_url, as part of the function's configuration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dense dim: 1536 (1536,)\n",
      "dense dim: 1536 (1536,)\n"
     ]
    }
   ],
   "source": [
    "# initialize using api_key directly.\n",
    "openai_ef = model.dense.OpenAIEmbeddingFunction(model_name=\"text-embedding-3-small\", api_key=\"sk-api-key\")\n",
    "# get the embeddings\n",
    "queries_embeddings = openai_ef.encode_queries(queries)\n",
    "docs_embeddings = openai_ef.encode_documents(docs)\n",
    "print(\"dense dim:\", openai_ef.dim, queries_embeddings[0].shape)\n",
    "print(\"dense dim:\", openai_ef.dim, docs_embeddings[0].shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sentence Transformer Embedding Function\n",
    "\n",
    "In addition to hosted services like OpenAI, there exists a variety of powerful open-source dense embedding models. For these, the SentenceTransformerEmbeddingFunction based on [Sentence-Transformer](https://www.sbert.net/) can be utilized to extract text embeddings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dense dim: 384 (384,)\n",
      "dense dim: 384 (384,)\n"
     ]
    }
   ],
   "source": [
    "docs = [\n",
    "    \"Artificial intelligence was founded as an academic discipline in 1956.\",\n",
    "    \"Alan Turing was the first person to conduct substantial research in AI.\",\n",
    "    \"Born in Maida Vale, London, Turing was raised in southern England.\",\n",
    "    \"The Turing Test, proposed by Alan Turing, is a measure of a machine's ability to exhibit intelligent behavior.\",\n",
    "    \"Deep learning is a subset of machine learning in artificial intelligence that has networks capable of learning unsupervised from data that is unstructured or unlabeled.\",\n",
    "    \"The concept of neural networks, which are vital to deep learning algorithms, was inspired by the understanding of the human brain's structure and function.\",\n",
    "    \"Artificial intelligence applications range from natural language processing to expert systems, and from automated reasoning to machine learning.\",\n",
    "    \"The development of quantum computing holds the potential to drastically increase the processing power available for artificial intelligence systems.\",\n",
    "    \"In the field of robotics, artificial intelligence is used to enable autonomous decision-making by robots in complex environments.\",\n",
    "    \"Ethical considerations in AI research and application are becoming increasingly important as the technology advances and becomes more integrated into daily life.\",\n",
    "    \"Reinforcement learning, a type of machine learning algorithm, enables an agent to learn in an interactive environment by trial and error using feedback from its own actions and experiences.\",\n",
    "    \"AI has the potential to revolutionize industries by optimizing processes, enhancing decision-making, and creating new opportunities for innovation.\"\n",
    "]\n",
    "\n",
    "from pymilvus import model\n",
    "\n",
    "# initialize the SentenceTransformerEmbeddingFunction\n",
    "sentence_transformer_ef = model.dense.SentenceTransformerEmbeddingFunction(\n",
    "    model_name=\"all-MiniLM-L6-v2\", # Specify the model name\n",
    "    device=\"cpu\" # Specify the device to use, e.g., 'cpu' or 'cuda:0'\n",
    ")\n",
    "\n",
    "# get the embeddings in general way.\n",
    "queries = docs\n",
    "queries_embeddings = sentence_transformer_ef(queries)\n",
    "docs_embeddings = sentence_transformer_ef(docs)\n",
    "\n",
    "# get the embeddings in specified way.\n",
    "queries_embeddings = sentence_transformer_ef.encode_queries(queries)\n",
    "docs_embeddings = sentence_transformer_ef.encode_documents(docs)\n",
    "\n",
    "print(\"dense dim:\", sentence_transformer_ef.dim, queries_embeddings[0].shape)\n",
    "print(\"dense dim:\", sentence_transformer_ef.dim, docs_embeddings[0].shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "when using SentenceTransformerEmbeddingFunction, `encoding_queries` and `encoding_documents` would prepend **query_instruction** and **doc_instruction** respectively, others are same.  \n",
    "- `sentence_transformer_ef(texts)`: prepend nothing, just process the raw text.\n",
    "- `sentence_transformer_ef.encode_queries(queries)`: prepend the **query_instruction** to each query. \n",
    "- `sentence_transformer_ef.encode_documents(documents)`: prepend the **doc_instruction** to each document.\n",
    "\n",
    "Additionally, the initialization of SentenceTransformerEmbeddingFunction may incorporate features from Sentence Transformer, such as specifying parameters like the `batch_size`. Some models require adding an instruction before the actual text input. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dense dim: 384 (384,)\n",
      "dense dim: 384 (384,)\n"
     ]
    }
   ],
   "source": [
    "#BAAI/bge-small-en-v1.5 suggest prepend a instruction when generate embedding.\n",
    "sentence_transformer_ef = model.dense.SentenceTransformerEmbeddingFunction(\n",
    "    model_name=\"BAAI/bge-small-en-v1.5\",\n",
    "    device=\"cpu\",\n",
    "    batch_size=8,\n",
    "    query_instruction=\"Represent this sentence for searching relevant passages:\",\n",
    "    doc_instruction=\"Represent this sentence for searching relevant passages:\",\n",
    ")\n",
    "queries_embeddings = sentence_transformer_ef.encode_queries(docs)\n",
    "docs_embeddings = sentence_transformer_ef.encode_documents(docs)\n",
    "\n",
    "print(\"dense dim:\", sentence_transformer_ef.dim, queries_embeddings[0].shape)\n",
    "print(\"dense dim:\", sentence_transformer_ef.dim, docs_embeddings[0].shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Sparse embedding** represents words or phrases using vectors where most elements are zero, with only one non-zero element indicating the presence of a specific word in a vocabulary. Sparse embeddings models are efficient and interpretable, making them suitable for tasks where exact term matches are crucial\n",
    "\n",
    "## Splade Embedding Function\n",
    "\n",
    "[SPLADE](https://arxiv.org/abs/2109.10086) embedding is a model that offers highly sparse representations for documents and queries, inheriting desirable properties from bag-of-words (BOW) models such as exact term matching and efficiency. We can use SPLADE model easily with SpladeEmbeddingFunction.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sparse dim: 30522 (1, 30522)\n",
      "sparse dim: 30522 (1, 30522)\n"
     ]
    }
   ],
   "source": [
    "docs = [\n",
    "    \"Artificial intelligence was founded as an academic discipline in 1956.\",\n",
    "    \"Alan Turing was the first person to conduct substantial research in AI.\",\n",
    "    \"Born in Maida Vale, London, Turing was raised in southern England.\",\n",
    "    \"The Turing Test, proposed by Alan Turing, is a measure of a machine's ability to exhibit intelligent behavior.\",\n",
    "]\n",
    "  \n",
    "from pymilvus.model.sparse import SpladeEmbeddingFunction\n",
    "\n",
    "# default using model_name naver/splade-cocondenser-ensembledistil. \n",
    "# other valid options: \n",
    "# - naver/splade_v2_max\n",
    "# - naver/splade_v2_distil\n",
    "# - naver/splade-cocondenser-selfdistil.\n",
    "splade_ef = SpladeEmbeddingFunction()\n",
    "\n",
    "queries = docs\n",
    "\n",
    "# get the embeddings in general way.\n",
    "queries_embeddings = splade_ef(queries)\n",
    "docs_embeddings = splade_ef(docs)\n",
    "\n",
    "# get the embeddings in specified way.\n",
    "queries_embeddings = splade_ef.encode_queries(queries)\n",
    "docs_embeddings = splade_ef.encode_documents(docs)\n",
    "\n",
    "# since the output embeddings are in a 2D csr_array format, we convert them to a list for easier manipulation.\n",
    "print(\"sparse dim:\", splade_ef.dim, list(queries_embeddings)[0].shape)\n",
    "print(\"sparse dim:\", splade_ef.dim, list(docs_embeddings)[0].shape)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "when using SpladeEmbeddingFunbction, `encoding_queries` and `encoding_documents` would prepend **query_instruction** and **doc_instruction** respectively, and **k_tokens_query** is used to prune the query results. **k_tokens_document** is used to prune the document results.  \n",
    "- `splade_ef(texts)`: prepend nothing, it does not prune the results.\n",
    "- `splade_ef.encode_queries(queries)`: prepend the **query_instruction** to each query, **k_tokens_query** is used to prune the query results.\n",
    "- `splade_ef.encode_documents(documents)`: prepend the **doc_instruction** to each document, **k_tokens_document** is used to prune the document results. \n",
    "\n",
    "\n",
    "By default, the model outputs the result directly. However, there are situations where users may only want to retain the top k largest values for a desired valid value. In such cases, users can specify the parameters 'k_tokens_query' and 'k_tokens_document' for queries and documents, respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sparse dim: 30522 (1, 30522)\n",
      "sparse dim: 30522 (1, 30522)\n",
      "query embedding non zero elemments: 64\n",
      "document embedding non zero elemments: 128\n"
     ]
    }
   ],
   "source": [
    "#Initialize the SpladeEmbeddingFunction retaining the top 64 tokens for queries and the top 128 tokens for documents.\n",
    "splade_ef = SpladeEmbeddingFunction(device=\"cpu\", k_tokens_query=64, k_tokens_document=128)\n",
    "    \n",
    "queries_embeddings = splade_ef.encode_queries(queries)\n",
    "docs_embeddings = splade_ef.encode_documents(docs)\n",
    "\n",
    "print(\"sparse dim:\", splade_ef.dim, list(queries_embeddings)[0].shape)\n",
    "print(\"sparse dim:\", splade_ef.dim, list(docs_embeddings)[0].shape)\n",
    "\n",
    "print(\"query embedding non zero elemments:\", list(queries_embeddings)[0].nnz)\n",
    "print(\"document embedding non zero elemments:\", list(docs_embeddings)[0].nnz)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## BM25 Embedding Function\n",
    "[BM25](https://en.wikipedia.org/wiki/Okapi_BM25) is a ranking function used in information retrieval to estimate the relevance of documents to a given search query. It enhances the basic term frequency approach by incorporating document length normalization and term frequency saturation. BM25 can generate sparse embeddings by representing documents as vectors of term importance scores, allowing for efficient retrieval and ranking in sparse vector spaces.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tokens: ['artifici', 'intellig', 'found', 'academ', 'disciplin', '1956']\n"
     ]
    }
   ],
   "source": [
    "from pymilvus.model.sparse.bm25.tokenizers import build_default_analyzer\n",
    "from pymilvus.model.sparse import BM25EmbeddingFunction\n",
    "\n",
    "# there are some built-in analyzers for several languages, now we use 'en' for English.\n",
    "analyzer = build_default_analyzer(language=\"en\")\n",
    "\n",
    "corpus = [\n",
    "    \"Artificial intelligence was founded as an academic discipline in 1956.\",\n",
    "    \"Alan Turing was the first person to conduct substantial research in AI.\",\n",
    "    \"Born in Maida Vale, London, Turing was raised in southern England.\",\n",
    "]\n",
    "\n",
    "# analyzer can tokenize the text into tokens\n",
    "tokens = analyzer(corpus[0])\n",
    "print(\"tokens:\", tokens)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The BM25 algorithm processes text by first breaking it into tokens using a built-in analyzer, as shown with English language tokens like 'artifici', 'intellig', and 'academ'. It then gathers statistics on these tokens, evaluating their frequency and distribution across documents. The core of BM25 calculates the relevance score of each token based on its importance, with rarer tokens receiving higher scores. This concise process enables effective ranking of documents by relevance to a query. \n",
    "\n",
    " So we need to fit a dataset(or corpus) to get the statistics.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sparse dim: 21 (1, 21)\n",
      "sparse dim: 21 (1, 21)\n"
     ]
    }
   ],
   "source": [
    "bm25_ef = BM25EmbeddingFunction(analyzer)\n",
    "\n",
    "# Fit the model on the corpus to get the statstics of the corpus.\n",
    "bm25_ef.fit(corpus)\n",
    "docs = [\n",
    "    \"The field of artificial intelligence was established as an academic subject in 1956.\",\n",
    "    \"Alan Turing was the pioneer in conducting significant research in artificial intelligence.\",\n",
    "    \"Originating in Maida Vale, London, Turing grew up in the southern regions of England.\",\n",
    "    \"In 1956, artificial intelligence emerged as a scholarly field.\",\n",
    "    \"Turing, originally from Maida Vale, London, was brought up in the south of England.\"\n",
    "]\n",
    "queries = docs\n",
    "\n",
    "# get the embeddings in specified way.\n",
    "queries_embeddings = bm25_ef.encode_queries(queries)\n",
    "docs_embeddings = bm25_ef.encode_documents(docs)\n",
    "\n",
    "# Since the output embeddings are in a 2D csr_array format, we convert them to a list for easier manipulation.\n",
    "print(\"sparse dim:\", bm25_ef.dim, list(queries_embeddings)[0].shape)\n",
    "print(\"sparse dim:\", bm25_ef.dim, list(docs_embeddings)[0].shape)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "when using BM25EmbeddingFunction, `encoding_queries` and `encoding_documents` are not exchangable mathematically. There is no avaliable `bm25_ef(texts)` implmented.  \n",
    "- `bm25_ef(texts)`: is not avaliable.\n",
    "- `bm25_ef.encode_queries(queries)`: has its distinct implmentation.\n",
    "- `bm25_ef.encode_documents(documents)`: has its distinct implmentation.\n",
    "\n",
    "Fitting data each time is time-consuming; we offer save and load features for efficiency."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sparse dim: 21 (1, 21)\n",
      "sparse dim: 21 (1, 21)\n"
     ]
    }
   ],
   "source": [
    "bm25_ef.save(\"example.json\")\n",
    "new_bm25_ef = BM25EmbeddingFunction(analyzer)\n",
    "new_bm25_ef.load(\"example.json\")\n",
    "\n",
    "queries_embeddings = new_bm25_ef.encode_queries(queries)\n",
    "docs_embeddings = new_bm25_ef.encode_documents(docs)\n",
    "\n",
    "# Since the output embeddings are in a 2D csr_array format, we convert them to a list for easier manipulation.\n",
    "print(\"sparse dim:\", new_bm25_ef.dim, list(queries_embeddings)[0].shape)\n",
    "print(\"sparse dim:\", new_bm25_ef.dim, list(docs_embeddings)[0].shape)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "Calculating statistics within similar distributions is crucial for achieving accurate results. But when there lacks corpus, we provide pre-built, fitted data on the MS MARCO dataset specifically for English. This will download a pre-built json file into your local path."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "path is None, using default bm25_msmarco_v1.json.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sparse dim: 3889385 (1, 3889385)\n",
      "sparse dim: 3889385 (1, 3889385)\n"
     ]
    }
   ],
   "source": [
    "prebuilt_bm25_ef = BM25EmbeddingFunction(analyzer)\n",
    "# load the pre-built json file without fitting the corpus.\n",
    "prebuilt_bm25_ef.load()\n",
    "queries_embeddings = prebuilt_bm25_ef.encode_queries(queries)\n",
    "docs_embeddings = prebuilt_bm25_ef.encode_documents(docs)\n",
    "# Since the output embeddings are in a 2D csr_array format, we convert them to a list for easier manipulation.\n",
    "print(\"sparse dim:\", prebuilt_bm25_ef.dim, list(queries_embeddings)[0].shape)\n",
    "print(\"sparse dim:\", prebuilt_bm25_ef.dim, list(docs_embeddings)[0].shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, let's show analyzers for other languages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Building prefix dict from the default dictionary ...\n",
      "Loading model from cache /tmp/jieba.cache\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading model cost 0.627 seconds.\n",
      "Prefix dict has been built successfully.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "de_tokens ['alan', 'turing', 'erst', 'person', 'umfangreich', 'forschung', 'bereich', 'ki', 'durchfuhrt']\n",
      "fr_tokens ['alan', 'turing', 'premi', 'person', 'men', 'recherch', 'approfond', 'ia']\n",
      "zh_tokens ['艾伦', '图灵', '第一个', '人工智能', '领域', '深入研究', '人']\n"
     ]
    }
   ],
   "source": [
    "de_text = \"Alan Turing war die erste Person, die umfangreiche Forschungen im Bereich der KI durchführte.\"\n",
    "fr_text = \"Alan Turing était la première personne à mener des recherches approfondies en IA.\"\n",
    "zh_text = \"艾伦·图灵是第一个进行人工智能领域深入研究的人。\"\n",
    "\n",
    "de_analyzer = build_default_analyzer(language=\"de\")\n",
    "fr_analyzer = build_default_analyzer(language=\"fr\")\n",
    "zh_analyzer = build_default_analyzer(language=\"zh\")\n",
    "\n",
    "de_tokens = de_analyzer(de_text)\n",
    "fr_tokens = fr_analyzer(fr_text)\n",
    "zh_tokens = zh_analyzer(zh_text)\n",
    "\n",
    "print(\"de_tokens\", de_tokens)\n",
    "print(\"fr_tokens\", fr_tokens)\n",
    "print(\"zh_tokens\", zh_tokens)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "In the realm of embedding models, hybrid architectures exist that generate both dense and sparse embeddings. We refer to these models as hybrid models and introduce BGE-M3 as an example of such a model.\n",
    "## BGE-M3 Embedding Function\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The [BGE-M3](https://arxiv.org/abs/2402.03216) is named for its capabilities in Multi-Linguality, Multi-Functionality, and Multi-Granularity. Capable of supporting over 100 languages, BGE-M3 sets new benchmarks in multi-lingual and cross-lingual retrieval tasks. Its unique ability to perform dense retrieval, multi-vector retrieval, and sparse retrieval within a single framework makes it an ideal choice for a wide range of information retrieval (IR) applications.  \n",
    "\n",
    "(caution: when running this section, please restart the jupyter python kernel firstly.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cb3aa297af2e4224bb3f65f7a8ed69f5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Fetching 23 files:   0%|          | 0/23 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading existing colbert_linear and sparse_linear---------\n",
      "dense query dim: 1024 (1024,)\n",
      "dense document dim: 1024 (1024,)\n",
      "sparse query dim: 250002 (1, 250002)\n",
      "sparse document dim: 250002 (1, 250002)\n"
     ]
    }
   ],
   "source": [
    "docs = [\n",
    "    \"Artificial intelligence was founded as an academic discipline in 1956.\",\n",
    "    \"Alan Turing was the first person to conduct substantial research in AI.\",\n",
    "    \"Born in Maida Vale, London, Turing was raised in southern England.\",\n",
    "    \"The Turing Test, proposed by Alan Turing, is a measure of a machine's ability to exhibit intelligent behavior.\",\n",
    "]\n",
    "  \n",
    "from pymilvus.model.hybrid import BGEM3EmbeddingFunction\n",
    "import numpy as np\n",
    "\n",
    "# please set the use_fp16 to False when you are using cpu.\n",
    "# by default the return options is:\n",
    "#  return_dense True\n",
    "#  return_sparse True\n",
    "#  return_colbert_vecs False \n",
    "bge_m3_ef = BGEM3EmbeddingFunction(\n",
    "    model_name='BAAI/bge-m3', # Specify the model name\n",
    "    device='cpu', # Specify the device to use, e.g., 'cpu' or 'cuda:0'\n",
    "    use_fp16=False # Specify whether to use fp16. Set to `False` if `device` is `cpu`.\n",
    ")\n",
    "queries = docs\n",
    "\n",
    "# get the embeddings in general way.\n",
    "queries_embeddings = bge_m3_ef(queries)\n",
    "docs_embeddings = bge_m3_ef(docs)\n",
    "\n",
    "# get the embeddings in specified way.\n",
    "queries_embeddings = bge_m3_ef.encode_queries(queries)\n",
    "docs_embeddings = bge_m3_ef.encode_documents(docs)\n",
    "\n",
    "print(\"dense query dim:\", bge_m3_ef.dim[\"dense\"], queries_embeddings[\"dense\"][0].shape)\n",
    "print(\"dense document dim:\", bge_m3_ef.dim[\"dense\"], docs_embeddings[\"dense\"][0].shape)\n",
    "\n",
    "# Since the sparse embeddings are in a 2D csr_array format, we convert them to a list for easier manipulation.\n",
    "print(\"sparse query dim:\", bge_m3_ef.dim[\"sparse\"], list(queries_embeddings[\"sparse\"])[0].shape)\n",
    "print(\"sparse document dim:\", bge_m3_ef.dim[\"sparse\"], list(docs_embeddings[\"sparse\"])[0].shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "when using BGEM3EmbeddingFunction, `encoding_queries` and `encoding_documents` are exactly same procedure, one can use `bge_m3_ef(texts)` instead.  \n",
    "- `bge_m3_ef(texts)`: same with other two functions.\n",
    "- `bge_m3_ef.encode_queries(queries)`: same with other two functions.\n",
    "- `bge_m3_ef.encode_documents(documents)`: same with other two functions.\n",
    "\n",
    "\n",
    "Although BGE-M3 is capable of generating both dense and sparse embeddings simultaneously, it can also be configured to function as a standard dense or sparse embedding generator by adjusting its return options."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2c759af4be9f4dc8bbf669a68163c88b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Fetching 23 files:   0%|          | 0/23 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading existing colbert_linear and sparse_linear---------\n",
      "dense query dim: 1024 (1024,)\n"
     ]
    }
   ],
   "source": [
    "# use bge-m3 as a dense embedding.\n",
    "bge_m3_ef = BGEM3EmbeddingFunction(\n",
    "    model_name='BAAI/bge-m3', # Specify the model name\n",
    "    device='cpu', # Specify the device to use, e.g., 'cpu' or 'cuda:0'\n",
    "    use_fp16=False, # Specify whether to use fp16. Set to `False` if `device` is `cpu`.\n",
    "    return_sparse=False # only allow the dense embedding output\n",
    ")\n",
    "queries = docs\n",
    "docs_embeddings = bge_m3_ef.encode_documents(docs)\n",
    "print(\"dense query dim:\", bge_m3_ef.dim[\"dense\"], queries_embeddings[\"dense\"][0].shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "Similarly, we can utilize the BGEM3EmbeddingFunction as a sparse embedding function.\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "414faea32b3943519be0cd94b93e5924",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Fetching 23 files:   0%|          | 0/23 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading existing colbert_linear and sparse_linear---------\n",
      "sparse query dim: 250002 (1, 250002)\n"
     ]
    }
   ],
   "source": [
    "# use bge-m3 as a sparse embedding.\n",
    "bge_m3_ef = BGEM3EmbeddingFunction(\n",
    "    model_name='BAAI/bge-m3', # Specify the model name\n",
    "    device='cpu', # Specify the device to use, e.g., 'cpu' or 'cuda:0'\n",
    "    use_fp16=False, # Specify whether to use fp16. Set to `False` if `device` is `cpu`.\n",
    "    return_dense=False # only allow the sparse embedding output\n",
    ")\n",
    "queries = docs\n",
    "docs_embeddings = bge_m3_ef.encode_documents(docs)\n",
    "# Since the sparse embeddings are in a 2D csr_array format, we convert them to a list for easier manipulation.\n",
    "print(\"sparse query dim:\", bge_m3_ef.dim[\"sparse\"], list(queries_embeddings[\"sparse\"])[0].shape)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
